package com.lf.mygpt.service.openai;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import com.unfbx.chatgpt.sse.ConsoleEventSourceListener;
import jdk.nashorn.internal.objects.annotations.Constructor;
import lombok.Getter;
import lombok.Setter;
import okhttp3.Response;
import okhttp3.ResponseBody;
import okhttp3.sse.EventSource;
import okhttp3.sse.EventSourceListener;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;
import java.util.Objects;

@Getter
@Setter
// 事件监听器类，用于处理 OpenAI 接口的流式返回，使用 SSE 的方式流式返回给调用方
public class OpenAiStreamListener extends EventSourceListener {
    // 引入日志组件
    private static final Logger log = LoggerFactory.getLogger(ConsoleEventSourceListener.class);
    // SSE 结构
    private SseEmitter sseEmitter;
    // 用于存储最终返回的结果
    private StringBuilder result = new StringBuilder();

    public OpenAiStreamListener(SseEmitter sseEmitter) {
        this.sseEmitter = sseEmitter;
    }

    @Override
    public void onOpen(EventSource eventSource, Response response) {
        log.info("OpenAI建立sse连接...");
    }

    @Override
    public void onEvent(EventSource eventSource, String id, String type, String data) {
        log.info("OpenAI返回数据：{}", data);
        if (data.equals("[DONE]")) {
            log.info("OpenAI返回数据结束了");
            sseEmitter.complete();
        } else {
            try {
                sseEmitter.send(data);
                JSONObject dataJson = JSON.parseObject(data, JSONObject.class);
                if (dataJson.containsKey("choices")) {
                    JSONObject choice = (JSONObject) dataJson.getJSONArray("choices").get(0);
                    String content = choice.getJSONObject("delta").getString("content");
                    if (content != null) {
                        result.append(content);
                    }
                }
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
    }

    @Override
    public void onClosed(EventSource eventSource) {
        log.info("OpenAI关闭sse连接...");
        sseEmitter.complete();
    }

    @Override
    public void onFailure(EventSource eventSource, Throwable t, Response response) {
        try {
            if (Objects.isNull(response)) {
                log.error("OpenAI  sse连接异常:{}", t);
                eventSource.cancel();
            } else {
                ResponseBody body = response.body();
                if (Objects.nonNull(body)) {
                    log.error("OpenAI  sse连接异常data：{}，异常：{}", body.string(), t);
                } else {
                    log.error("OpenAI  sse连接异常data：{}，异常：{}", response, t);
                }
                eventSource.cancel();
            }
            sseEmitter.complete();
        } catch (Exception e) {
            e.printStackTrace();
        }
    }
}
